package tencent.music;

import java.util.HashSet;
import java.util.LinkedList;
import java.util.Queue;
import java.util.Set;

/**
 * @author habitplus
 * @problem
 * @desc
 * @link
 * @since 2021-08-26 18:47
 */
public class T20210826N1 {
    public TreeNode pruneLeaves(TreeNode root) {
        // write code here
        if (root == null) {
            return root;
        }

        Queue<TreeNode> que = new LinkedList<>();
        TreeNode t;

        Set<TreeNode> set = new HashSet<>();
        que.offer(root);
        // 记录所有叶子结点
        while (!que.isEmpty()) {
            for (int i = que.size(); i > 0; i--) {
                t = que.poll();
                if (t.left == null && t.right == null) {
                    // t 为叶子结点
                    set.add(t);
                } else {
                    if (t.left != null) que.offer(t.left);
                    if (t.right != null) que.add(t.right);
                }
            }
        }

        Set<TreeNode> delSet = new HashSet<>();
        que.offer(root);
        // 记录要删除的节点
        while (!que.isEmpty()) {
            for (int i = que.size(); i > 0; i--) {
                t = que.poll();

                if (set.contains(t.left) || set.contains(t.right)) {
                    delSet.add(t);
                    continue;
                }

                if (t.left != null) que.offer(t.left);
                if (t.right != null) que.offer(t.right);

            }
        }

        que.offer(root);
        // 删除节点
        while (!que.isEmpty()) {
            for (int i = que.size(); i > 0; i--) {
                t = que.poll();

                if (t.left != null) {
                    if (delSet.contains(t.left)) {
                        t.left = null;
                    } else {
                        que.offer(t.left);
                    }
                }
                if (t.right != null) {
                    if (delSet.contains(t.right)) {
                        t.right = null;
                    } else {
                        que.offer(t.right);
                    }
                }

            }
        }
        return root;
    }
}

class TreeNode {
    int val = 0;
    TreeNode left = null;
    TreeNode right = null;

    public TreeNode(int val) {
        this.val = val;
    }
}
